import numpy as np
from threading import Event, Thread, Lock
import time
from typing import List, Union, Callable

from ezrobot.utils.general import accurate_sleep, CircularBuffer


class BaseSensor(object):

    def recv_data(self, **kwargs):
        raise NotImplementedError("This method should be implemented in the subclass.")


class SensorManager:

    def __init__(
        self, 
        sensor: Union[BaseSensor, List[BaseSensor]],
        frequency = 30,
        buffer_size: int = 10,
        streaming: bool = False,
        callback: Callable = None
    ):
        """
        Base class for sensors
        """
        self._buffer = CircularBuffer(buffer_size)
        self._sleep_time_s = 1 / frequency
        self._is_streaming = Event()
        self._streaming_thread = None
        # self._lock = Lock()
        self._sensor = sensor
        self._callback = callback
        if streaming:
            self.start_streaming()
    
    def start_streaming(self):
        """
        Start streaming data from the sensor.
        
        Parameters:
        - delay_time: float, optional, default: 0.0, the delay time before collecting data.
        """
        if self._is_streaming.is_set():
            return
        self._streaming_thread = Thread(target=self._streaming, daemon=True)
        self._streaming_thread.start()

    def _streaming(self):
        self._is_streaming.set()
        while self._is_streaming.is_set():
            t0 = time.time()

            if isinstance(self._sensor, list):
                self._buffer.put({f"{sensor.__class__.__name__}_{str(sensor.index)}":sensor.recv_data() for sensor in self._sensor})
                # self._buffer.put([sensor.recv_data() for sensor in self._sensor])
            else:
                self._buffer.put(self._sensor.recv_data())
            
            # callback
            if self._callback is not None:
                self._callback(self._buffer.get())

            elapsed_time = time.time() - t0
            if self._sleep_time_s > elapsed_time:
                accurate_sleep(self._sleep_time_s - elapsed_time)

    def stop_streaming(self):
        """
        Stop streaming data from the sensor.
        """
        self._is_streaming.clear()
        if self._streaming_thread is not None:
            self._streaming_thread.join()
            self._streaming_thread = None
        self._buffer.clear()

    def get_data(self, **kwargs):
        if not self._is_streaming.is_set():
            raise RuntimeError("Streaming is not started. Please start streaming before getting data.")
        data = self._buffer.get()
        if data:
            return data
        else:
            accurate_sleep(0.05)
            return self.get_data(**kwargs)

if __name__ == "__main__":
    # test multiple sensors
    from ezrobot.sensor.encoder import AngleEncoder
    from ezrobot.sensor.vive_tracker import ViveTracker
    from ezrobot.sensor.tactile_sensor import TactileSensor
    from ezrobot.sensor.realsense import RealSense

    encoder = AngleEncoder("COM14", 1)
    tracker = ViveTracker(1)
    tactile_sensor = TactileSensor("OG000086")
    tactile_sensor1 = TactileSensor("OG000131")
    tactile_sensor2 = TactileSensor("OG000053")
    rs = RealSense()
    sensor = SensorManager(sensor=[encoder, tracker, tactile_sensor, tactile_sensor1, tactile_sensor2, rs], frequency=30, buffer_size=10, streaming=True)
    # sensor = SensorManager(sensor=[encoder, tracker], frequency=30, buffer_size=10, streaming=True)
    i= 0
    while True:
       print(i)
       i+=1
       data = sensor.get_data()
       print(data['RealSense_021422060263']['color'].shape)
       print(data['RealSense_021422060263']['depth'].shape)
       accurate_sleep(0.01)